#pragma once 
#include <iostream>
#include <string>
using namespace std;

enum Color{RED, BLACK};

template<class T>
struct RBTreeNode 
{
  RBTreeNode(const T& data)
    :_left(nullptr)
    ,_right(nullptr)
    ,_parent(nullptr)
    ,_col(RED)
    ,_data(data)
  {}
  RBTreeNode<T>* _left;
  RBTreeNode<T>* _right;
  RBTreeNode<T>* _parent;
  Color _col;
  T _data;
};




template<class T, class Ref, class Ptr>
struct rbtree_iterator
{
  typedef RBTreeNode<T> Node;
  typedef rbtree_iterator<T, Ref, Ptr> iterator;
  rbtree_iterator(Node* node)
    :_node(node)
  {}
  iterator operator++();
  iterator operator--();
  bool operator!=(const iterator& it) { return _node != it._node; }
  Ref operator*() { return _node->_data; }
  Ptr operator->() { return &_node->_data;  }
  Node* _node;
};

template<class T, class Ref, class Ptr>
rbtree_iterator<T, Ref, Ptr> rbtree_iterator<T, Ref, Ptr>::operator++()
{
  Node* cur = _node;
  if (cur == nullptr)
  {
    _node = nullptr;
    return *this;
  }
  if (cur->_right)
  {
    cur = cur->_right;
    while (cur->_left)
    {
      cur = cur->_left;
    }
    _node = cur;
    return *this;
  }
  else 
  {
    Node* parent = cur->_parent;
    while (parent && parent->_right == cur)
    {
      cur = parent;
      parent = parent->_parent;
    }
    _node = parent;
    return *this;
  }
}

template<class T, class Ref, class Ptr>
rbtree_iterator<T, Ref, Ptr> rbtree_iterator<T, Ref, Ptr>::operator--()
{
  Node* cur = _node;
  if (cur == nullptr)
  {
    return *this;
  }
  if (_node->_left)
  {
    cur = _node->_left;
    while (cur->_right)
    {
      cur = cur->_right;
    }
    _node = cur;
    return *this;
  }
  else 
  {
    Node* parent = cur->_parent;
    while (parent && parent->_left == cur)
    {
      cur = parent;
      parent = parent->_parent;
    }
    _node = parent;
    return *this;
  }
}

template<class K, class T, class KeyOfT>
struct RBTree 
{
  typedef rbtree_iterator<T, T&, T*> iterator;
  typedef RBTreeNode<T> Node;
  RBTree()
    :_root(nullptr)
  {}
  ~RBTree() { Destroy(_root); };
  RBTree(const RBTree<K, T, KeyOfT>& rb) { _root = Copy(rb._root) ; }
  RBTree& operator=(RBTree tree) {swap(_root, tree._root); return *this;}
  iterator begin();
  iterator end();
  iterator find(const K& key);
  pair<iterator, bool> insert(const T& data);
  private:
  void Lrotate(Node* cur);
  void Rrotate(Node* cur);
  void LRrotate(Node* cur);
  void RLrotate(Node* cur);
  void Destroy(Node* root);
  Node* Copy(Node* root);
  Node* _root;
};

template<class K, class T, class KeyOfT>
RBTreeNode<T>* RBTree<K, T, KeyOfT>::Copy(Node* root)
{
  if (root == nullptr)
    return nullptr;

  Node* newRoot = new Node(root->_data);
  newRoot->_left = Copy(root->_left);
  if (newRoot->_left)
    newRoot->_left->_parent = newRoot;
  newRoot->_right = Copy(root->_right);
  if (newRoot->_right)
    newRoot->_right->_parent = newRoot;
  return newRoot;
}
template<class K, class T, class KeyOfT>
void RBTree<K, T, KeyOfT>::Destroy(Node* root)
{
  if (root == nullptr)
  {
    return;
  }
  else 
  {
    Destroy(root->_left);
    Destroy(root->_right);
    delete root;
  }

}

template<class K, class T, class KeyOfT>
typename RBTree<K, T, KeyOfT>::iterator RBTree<K, T, KeyOfT>::begin()
{
  Node* min = _root;
  while (min && min->_left)
  {
    min = min->_left;
  }
  return iterator(min);

}
template<class K, class T, class KeyOfT>
typename RBTree<K, T, KeyOfT>::iterator RBTree<K, T, KeyOfT>::end()
{
  return iterator(nullptr);
}

template<class K, class T, class KeyOfT>
typename RBTree<K, T, KeyOfT>::iterator RBTree<K, T, KeyOfT>::find(const K& key)
{
  KeyOfT kot;
  Node* cur = _root;
  while (cur)
  {
    if (key > kot(cur->_data))
    {
      cur = cur->_right;
    }
    else if (key < kot(cur->_data))
    {
      cur = cur->_left;
    }
    else 
    {
      return iterator(cur);
    }
  }
  return iterator(nullptr);
}
template<class K, class T, class KeyOfT>
pair<typename RBTree<K, T, KeyOfT>::iterator, bool> RBTree<K, T, KeyOfT>::insert(const T& data)
{
  Node* newnode;
  if (_root == nullptr)
		{
			_root = new Node(data);
			_root->_col = BLACK;
			return make_pair(iterator(_root), true);
		}
		else
		{
			KeyOfT kot;
			Node* cur = _root;
			Node* parent = nullptr;
			while (cur)
			{
				if (kot(data) > kot(cur->_data))
				{
					parent = cur;
					cur = cur->_right;
				}
				else if (kot(data) < kot(cur->_data))
				{
					parent = cur;
					cur = cur->_left;
				}
				else
				{
					return make_pair(iterator(cur), false);
				}
			}
			cur = new Node(data);
      newnode = cur;
			cur->_parent = parent;
			if (kot(cur->_data) > kot(parent->_data))
			{
				parent->_right = cur;
			}
			else
			{
				parent->_left = cur;
			}
			while (parent && parent->_col == RED)
			{
				Node* grandfather = parent->_parent;
				Node* uncle;
				if (parent == grandfather->_left)
				{
					uncle = grandfather->_right;
					if (uncle && uncle->_col == RED)
					{
						uncle->_col = parent->_col = BLACK;
						grandfather->_col = RED;
						cur = grandfather;
						parent = grandfather->_parent;
					}
					else if (uncle && uncle->_col == BLACK)
					{
						if (parent->_left == cur)
						{
							Rrotate(parent);
              _root->_col = BLACK;
							break;
						}
						else
						{
							//cur 是 parent的右节点
							LRrotate(cur);
              _root->_col = BLACK;
							break;
						}
					}
					else
					{
						//uncle 不存在
						if (parent->_left == cur)
						{
							Rrotate(parent);
              _root->_col = BLACK;
							break;
						}
						else
						{
							//cur 是 parent的右节点
							LRrotate(cur);
              _root->_col = BLACK;
							break;
						}
					}
				}
				else
				{
					uncle = grandfather->_left;
					if (uncle && uncle->_col == RED)
					{
						uncle->_col = parent->_col = BLACK;
						grandfather->_col = RED;
						cur = grandfather;
						parent = grandfather->_parent;
					}
					else if (uncle && uncle->_col == BLACK)
					{
						if (parent->_right == cur)
						{
							Lrotate(parent);
							break;
						}
						else
						{
							//cur 是 parent的左节点
							RLrotate(cur);
							break;
						}
					}
					else
					{
						//uncle 不存在
						if (parent->_right == cur)
						{
							Lrotate(parent);
							break;
						}
						else
						{
							//cur 是 parent的左节点
							RLrotate(cur);
							break;
						}
					}
				}
			}
		}
		     _root->_col = BLACK;
         return make_pair(iterator(newnode), true);
	}

template<class K, class T, class KeyOfT>
void RBTree<K, T, KeyOfT>::Lrotate(RBTreeNode<T>* cur)
{
	Node* parent = cur->_parent;
	Node* parpar = parent->_parent;
	Node* subL = cur->_left;
	parent->_right = subL;
	if (subL)
	{
		subL->_parent = parent;
	}
	cur->_parent = parpar;
	if (parpar)
	{
		if (parpar->_left == parent)
		{
			parpar->_left = cur;
		}
		else
		{
			parpar->_right = cur;
		}
	}
	else
	{
		_root = cur;
		cur->_parent = nullptr;
	}
	cur->_left = parent;
	parent->_parent = cur;
	cur->_col = BLACK;
	parent->_col = RED;

}
template<class K, class T, class KeyOfT>
void RBTree<K, T, KeyOfT>::Rrotate(RBTreeNode<T>* cur)
{
	Node* parent = cur->_parent;
	Node* parpar = parent->_parent;
	Node* subR = cur->_right;

	parent->_left = subR;
	if (subR)
	{
		subR->_parent = parent;
	}
	cur->_parent = parpar;
	if (parpar)
	{
		if (parpar->_left == parent)
		{
			parpar->_left = cur;
		}
		else
		{
			parpar->_right = cur;
		}

	}
	else
	{
		_root = cur;
		cur->_parent = nullptr;
	}
	cur->_right = parent;
	parent->_parent = cur;
	cur->_col = BLACK;
	parent->_col = RED;

}


template<class K, class T, class KeyOfT>
void RBTree<K, T, KeyOfT>::LRrotate(RBTreeNode<T>* cur)
{
  Lrotate(cur);
  Rrotate(cur);
}


template<class K, class T, class KeyOfT>
void RBTree<K, T, KeyOfT>::RLrotate(RBTreeNode<T>* cur)
{
  Rrotate(cur);
  Lrotate(cur);
}
